Huh, interesting. Likely need to extend ggml_mul_mat_id with more general alpha_i*A^B + beta_i*C and implement special casing in the kernels for alpha_i=0 to skip the compute.

Huh, interesting. Likely need to extend ggml_mul_mat_id with more general alpha_i*A^B + beta_i*C and implement special casing in the kernels for alpha_i=0 to skip the compute.

Hello, I'm also paying attention to the adaptation of the loncat-flash model. Suppose B is the expert weight. If we change mul_mat_ids to alpha * A^B + beta * C, since the weights of zero-compute experts are not saved and the experts activated for each token are different, there might be an interleaved expert activation scenario. In this case, we need to determine for each token whether the activated expert is a zero-compute expert. Would this have an impact on performance?

If we bypass the requirement of ggml_mul_mat_ids to arrange in token order and add a new operator that rearranges the activations in expert order, like torch._grouped_gemm, would this provide better adaptability? In this case, we would only need to determine whether the currently computed expert ID is valid or whether it is a zero-compute expert. See #18369

@hebangwen Not sure I follow, but I think simply setting the coefficient like this should work:

# normal expert
alpha_i = 1.0f
beta_i  = 0.0f

# zero-compute expert
alpha_i = 0.0f
beta_i  = 1.0f

And the matrix C is just the MoE input (i.e. x_t from the paper).

It can be simpler to explain the mul_mat_id in simple terms like this:

• A normal mul_mat(A, B) calculates A @ B = C, simple.
• Now, instead of having just a single B, we have B as a stack of multiple matrices: B = (B0, B1, B2, ..., Bn)
• For MoE, we want to mul_mat A with a subset of B, something like: A @ (B2, B8, ...)
• So mul_mat_id takes an extra param, the indexes of elements in B to be used: mul_mat_id(A, B, (2, 8)) --> (A @ B2, A @ B8)

As @ggerganov suggested, I imagine the mul_mat_id will now take an extra alpha, beta params:

idx = (2, 8)
alpha = (0.3, 0.0)
mul_mat_id(A, B, idx, alpha) --> (A @ B2 * 0.3, A @ B8 * 0.0)

In the example above, computation for A @ B8 will be skipped as its beta value is 0.0

However, one issue is that the router weight alpha is non-zero. Indeed, it depends on the top_k operation to sort out the activated experts. So, I think just adding a check like if (idx >= B->ne[2]) could be enough? So if B only has 4 experts and the experts ID=5 is accessed, then it's out-of-bound; we skip the mul_mat in such case.

However, yet another problem, even when the idea above is implemented: The output dim of mul_mat_id will not be the same as the input, and more importantly, there are also other ops like mul or gelu in-between FFN gate/up/down. We can resolve this by assuming that output of the skipped mul mat will be all 0.0, but I think that's not a generic solution.

For the calculation of the beta_i*C term, it will also be a bit tricky, as we need to filter out only the coefficients (aka router weights) that correspond to the zero-computing experts. Something like torch.where could be necessary, but that's another rabbit hole I think

Seems like quite more works than I initially thought, so I think we should re-consider if this worth implementing. Currently, only longcat-flash family using this technique, so it can be quite risky to too many infrastructure to support it.

Yes, seems more complicated. Let's reconsider later in case this architecture shows any promise.